{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "342d5a4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "import logging\n",
    "import time\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db49d7af",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2659548b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from docling.datamodel.base_models import InputFormat\n",
    "from docling.datamodel.pipeline_options import PdfPipelineOptions\n",
    "from docling.document_converter import DocumentConverter, PdfFormatOption\n",
    "from docling.utils.export import generate_multimodal_pages\n",
    "from docling.utils.utils import create_hash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99b0327e",
   "metadata": {},
   "outputs": [],
   "source": [
    "_log = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89bf4322",
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGE_RESOLUTION_SCALE = 2.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c3dd816",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    logging.basicConfig(level=logging.INFO)\n",
    "    input_doc_path = r'D:\\YXT_Project\\vector_db_practices\\docling\\10027.pdf'\n",
    "\n",
    "    # input_doc_path = Path(\"D:\\YXT_Project\\vector_db_practices\\docling\\\\10027.pdf\")\n",
    "    output_dir = Path(\"scratch\")\n",
    "\n",
    "    # Important: For operating with page images, we must keep them, otherwise the DocumentConverter\n",
    "    # will destroy them for cleaning up memory.\n",
    "    # This is done by setting AssembleOptions.images_scale, which also defines the scale of images.\n",
    "    # scale=1 correspond of a standard 72 DPI image\n",
    "    pipeline_options = PdfPipelineOptions()\n",
    "    pipeline_options.ocr_options.lang = [\"fr\", \"de\", \"es\", \"en\"]  # example of languages for EasyOCR\n",
    "    pipeline_options.images_scale = IMAGE_RESOLUTION_SCALE\n",
    "    pipeline_options.generate_page_images = True\n",
    "\n",
    "    doc_converter = DocumentConverter(\n",
    "        format_options={\n",
    "            InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)\n",
    "        }\n",
    "    )\n",
    "\n",
    "    start_time = time.time()\n",
    "\n",
    "    conv_res = doc_converter.convert(input_doc_path)\n",
    "\n",
    "    output_dir.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    rows = []\n",
    "    for (\n",
    "        content_text,\n",
    "        content_md,\n",
    "        content_dt,\n",
    "        page_cells,\n",
    "        page_segments,\n",
    "        page,\n",
    "    ) in generate_multimodal_pages(conv_res):\n",
    "        dpi = page._default_image_scale * 72\n",
    "\n",
    "        rows.append(\n",
    "            {\n",
    "                \"document\": conv_res.input.file.name,\n",
    "                \"hash\": conv_res.input.document_hash,\n",
    "                \"page_hash\": create_hash(\n",
    "                    conv_res.input.document_hash + \":\" + str(page.page_no - 1)\n",
    "                ),\n",
    "                \"image\": {\n",
    "                    \"width\": page.image.width,\n",
    "                    \"height\": page.image.height,\n",
    "                    \"bytes\": page.image.tobytes(),\n",
    "                },\n",
    "                \"cells\": page_cells,\n",
    "                \"contents\": content_text,\n",
    "                \"contents_md\": content_md,\n",
    "                \"contents_dt\": content_dt,\n",
    "                \"segments\": page_segments,\n",
    "                \"extra\": {\n",
    "                    \"page_num\": page.page_no + 1,\n",
    "                    \"width_in_points\": page.size.width,\n",
    "                    \"height_in_points\": page.size.height,\n",
    "                    \"dpi\": dpi,\n",
    "                },\n",
    "            }\n",
    "        )\n",
    "\n",
    "    # Generate one parquet from all documents\n",
    "    df_result = pd.json_normalize(rows)\n",
    "    now = datetime.datetime.now()\n",
    "    output_filename = output_dir / f\"multimodal_{now:%Y-%m-%d_%H%M%S}.parquet\"\n",
    "    df_result.to_parquet(output_filename)\n",
    "\n",
    "    end_time = time.time() - start_time\n",
    "\n",
    "    _log.info(\n",
    "        f\"Document converted and multimodal pages generated in {end_time:.2f} seconds.\"\n",
    "    )\n",
    "\n",
    "    # This block demonstrates how the file can be opened with the HF datasets library\n",
    "    # from datasets import Dataset\n",
    "    # from PIL import Image\n",
    "    # multimodal_df = pd.read_parquet(output_filename)\n",
    "\n",
    "    # # Convert pandas DataFrame to Hugging Face Dataset and load bytes into image\n",
    "    # dataset = Dataset.from_pandas(multimodal_df)\n",
    "    # def transforms(examples):\n",
    "    #     examples[\"image\"] = Image.frombytes('RGB', (examples[\"image.width\"], examples[\"image.height\"]), examples[\"image.bytes\"], 'raw')\n",
    "    #     return examples\n",
    "    # dataset = dataset.map(transforms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65049add",
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
